#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>
int main(int argc, char *argv[])
{
    int myid, numprocs,n;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank( MPI_COMM_WORLD, &myid);
    MPI_Comm_size(MPI_COMM_WORLD, &numprocs);
    MPI_Status status; 
    double *A, *B, *C;
    if (myid == 0) 
    {
        printf("Enter the matrix size: ");
        scanf("%d", &n);

        A = (double *)malloc(n * n * sizeof(double));
        B = (double *)malloc(n * n * sizeof(double));
        C = (double *)malloc(n * n * sizeof(double));

        printf("Enter matrix A:\n");
        for (int i = 0; i < n * n; i++) {
            scanf("%lf,", &A[i]);
        }

        printf("Enter matrix B:\n");
        for (int i = 0; i < n * n; i++) {
            scanf("%lf,", &B[i]);
        }
    }
    
    int per = n/numprocs; //每个进程分配的矩阵行数
    
    double *subA = (double*) malloc( per * n * sizeof(double));
    double *subC = (double*) malloc( n * n * sizeof(double));
      
    //把A分块发送给不同进程    
    MPI_Scatter(A, per*n, MPI_DOUBLE, subA, per*n, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    //把B完整地发送给不同进程
    MPI_Bcast(B, n*n, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    
    for (int i = 0; i < n ; i++)
    {
        for (int j = 0; j < n; j++)
        { 
            for (int p = 0; p < per; p++)
            {
              subC[i * n + j] += subA[i * n + p] * B[p * n + j];
            }
        }
     }
     
    MPI_Gather(subC, n*n, MPI_DOUBLE, C, n*n, MPI_DOUBLE, 0, MPI_COMM_WORLD); 
    
    if (myid == 0) 
    {
        printf("Result matrix:\n");   
        for (int i = 0; i < n; i++) 
        {
            for (int j = 0; j < n; j++) 
            {
                printf("%lf", C[i * n + j]);
            }
            printf("\n");
            free(A);
            free(B);
            free(C);
        }
    }
        
    free(subA);
    free(subC);
    
    MPI_Finalize();
    
    return 0;
}
